#include "line_extraction.h"
#include <Eigen/Dense>
#include <algorithm>
#include <iostream>

namespace slam {
namespace mapping {

Line::Line(const sensor::PointCloud &laser_data,
           const LineExtractionParams &params)
    : laser_data_(laser_data), params_(params) {
  points_info_.clear();
  for (const auto &item : laser_data_) {
    LinePointInfo point_info;
    point_info.range = std::sqrt(item.position[0] * item.position[0] +
                                 item.position[1] * item.position[1]);
    point_info.angle = std::atan2(item.position[1], item.position[0]);
    points_info_.emplace_back(point_info);
  }
}

void Line::SetLineParas(const double angle, const double radius,
                        const boost::array<double, 4> &covariance,
                        const boost::array<double, 2> &start,
                        const boost::array<double, 2> &end) {
  angle_ = angle;
  radius_ = radius;
  covariance_ = covariance;
  start_ = start;
  end_ = end;
}

Line::~Line() {}

double Line::GetAngle() const { return angle_; }

const boost::array<double, 4> &Line::GetCovariance() const {
  return covariance_;
}

const boost::array<double, 2> &Line::GetEnd() const { return end_; }

double Line::GetRadius() const { return radius_; }

const boost::array<double, 2> &Line::GetStart() const { return start_; }

double Line::DistToPoint(const sensor::RangefinderPoint &point) {
  double p_rad = std::sqrt(std::pow(point.position[0], 2) +
                           std::pow(point.position[1], 2));
  double p_ang = std::atan2(point.position[1], point.position[0]);
  return std::fabs(p_rad * std::cos(p_ang - angle_) - radius_);
}

double Line::Length() const {
  return std::sqrt(std::pow(start_[0] - end_[0], 2) +
                   std::pow(start_[1] - end_[1], 2));
}

unsigned int Line::NumPoints() const { return laser_data_.size(); }

void Line::ProjectEndpoints() {
  double s = -1.0 / std::tan(angle_);
  double b = radius_ / std::sin(angle_);
  double x = start_[0];
  double y = start_[1];
  start_[0] = (s * y + x - s * b) / (std::pow(s, 2) + 1);
  start_[1] = (std::pow(s, 2) * y + s * x + b) / (std::pow(s, 2) + 1);
  x = end_[0];
  y = end_[1];
  end_[0] = (s * y + x - s * b) / (std::pow(s, 2) + 1);
  end_[1] = (std::pow(s, 2) * y + s * x + b) / (std::pow(s, 2) + 1);
}

// Methods for endpoint line fitting
void Line::EndpointFit() {
  start_[0] = laser_data_.front().position[0];
  start_[1] = laser_data_.front().position[1];
  end_[0] = laser_data_.back().position[0];
  end_[1] = laser_data_.back().position[1];
  AngleFromEndpoints();
  RadiusFromEndpoints();
}

void Line::AngleFromEndpoints() {
  double slope;
  if (std::fabs(end_[0] - start_[0]) > 1e-9) {
    slope = (end_[1] - start_[1]) / (end_[0] - start_[0]);
    angle_ = PIToPI(std::atan(slope) + M_PI / 2);
  } else {
    angle_ = 0.0;
  }
}

void Line::RadiusFromEndpoints() {
  radius_ = start_[0] * std::cos(angle_) + start_[1] * std::sin(angle_);
  if (radius_ < 0) {
    radius_ = -radius_;
    angle_ = PIToPI(angle_ + M_PI);
  }
}

// Methods for least squares line fitting
void Line::LeastSqFit() {
  CalcPointCovariances();
  double prev_radius = 0.0;
  double prev_angle = 0.0;
  while (std::fabs(radius_ - prev_radius) > params_.least_sq_radius_thresh ||
         std::fabs(angle_ - prev_angle) > params_.least_sq_angle_thresh) {
    prev_radius = radius_;
    prev_angle = angle_;
    CalcPointScalarCovariances();
    RadiusFromLeastSq();
    AngleFromLeastSq();
  }
  CalcCovariance();
  ProjectEndpoints();
}

const sensor::PointCloud &Line::GetLaserData() const { return laser_data_; }

void Line::AngleFromLeastSq() {
  CalcPointParameters();
  angle_ += AngleIncrement();
}

double Line::AngleIncrement() {
  const std::vector<double> &a = p_params_.a;
  const std::vector<double> &ap = p_params_.ap;
  const std::vector<double> &app = p_params_.app;
  const std::vector<double> &b = p_params_.b;
  const std::vector<double> &bp = p_params_.bp;
  const std::vector<double> &bpp = p_params_.bpp;
  const std::vector<double> &c = p_params_.c;
  const std::vector<double> &s = p_params_.s;

  double numerator = 0;
  double denominator = 0;
  for (std::size_t i = 0; i < a.size(); ++i) {
    numerator += (b[i] * ap[i] - a[i] * bp[i]) / std::pow(b[i], 2);
    denominator += ((app[i] * b[i] - a[i] * bpp[i]) * b[i] -
                    2 * (ap[i] * b[i] - a[i] * bp[i]) * bp[i]) /
                   std::pow(b[i], 3);
  }
  return -(numerator / denominator);
}

void Line::CalcCovariance() {
  covariance_[0] = p_rr_;

  const std::vector<double> &a = p_params_.a;
  const std::vector<double> &ap = p_params_.ap;
  const std::vector<double> &app = p_params_.app;
  const std::vector<double> &b = p_params_.b;
  const std::vector<double> &bp = p_params_.bp;
  const std::vector<double> &bpp = p_params_.bpp;
  const std::vector<double> &c = p_params_.c;
  const std::vector<double> &s = p_params_.s;

  double G = 0;
  double A = 0;
  double B = 0;
  double r, phi;
  CHECK_EQ(laser_data_.size(), points_info_.size());
  for (int i = 0; i < laser_data_.size(); i++) {
    r = points_info_[i].range;
    phi = points_info_[i].angle;
    G += ((app[i] * b[i] - a[i] * bpp[i]) * b[i] -
          2 * (ap[i] * b[i] - a[i] * bp[i]) * bp[i]) /
         std::pow(b[i], 3);
    A += 2 * r * std::sin(angle_ - phi) / b[i];
    B += 4 * std::pow(r, 2) * std::pow(std::sin(angle_ - phi), 2) / b[i];
  }
  covariance_[1] = p_rr_ * A / G;
  covariance_[2] = covariance_[1];
  covariance_[3] = std::pow(1.0 / G, 2) * B;
}

void Line::CalcPointCovariances() {
  point_covs_.clear();
  double r, phi, var_r, var_phi;
  CHECK_EQ(laser_data_.size(), points_info_.size());
  for (int i = 0; i < laser_data_.size(); i++) {
    r = points_info_[i].range;
    phi = points_info_[i].angle;    // bearing
    var_r = params_.range_var;      // range variance
    var_phi = params_.bearing_var;  // bearing variance
    boost::array<double, 4> Q;
    Q[0] = std::pow(r, 2) * var_phi * std::pow(std::sin(phi), 2) +
           var_r * std::pow(std::cos(phi), 2);
    Q[1] = -std::pow(r, 2) * var_phi * std::sin(2 * phi) / 2.0 +
           var_r * std::sin(2 * phi) / 2.0;
    Q[2] = Q[1];
    Q[3] = std::pow(r, 2) * var_phi * std::pow(std::cos(phi), 2) +
           var_r * std::pow(std::sin(phi), 2);
    point_covs_.push_back(Q);
  }
}

void Line::CalcPointParameters() {
  p_params_.a.clear();
  p_params_.ap.clear();
  p_params_.app.clear();
  p_params_.b.clear();
  p_params_.bp.clear();
  p_params_.bpp.clear();
  p_params_.c.clear();
  p_params_.s.clear();

  double r, phi, var_r, var_phi;
  double a, ap, app, b, bp, bpp, c, s;
  CHECK_EQ(laser_data_.size(), points_info_.size());
  for (int i = 0; i < laser_data_.size(); i++) {
    r = points_info_[i].range;      // range
    phi = points_info_[i].angle;    // bearing
    var_r = params_.range_var;      // range variance
    var_phi = params_.bearing_var;  // bearing variance
    c = std::cos(angle_ - phi);
    s = std::sin(angle_ - phi);
    a = std::pow(r * c - radius_, 2);
    ap = -2 * r * s * (r * c - radius_);
    app = 2 * std::pow(r, 2) * std::pow(s, 2) - 2 * r * c * (r * c - radius_);
    b = var_r * std::pow(c, 2) + var_phi * std::pow(r, 2) * std::pow(s, 2);
    bp = 2 * (std::pow(r, 2) * var_phi - var_r) * c * s;
    bpp = 2 * (std::pow(r, 2) * var_phi - var_r) *
          (std::pow(c, 2) - std::pow(s, 2));
    p_params_.a.push_back(a);
    p_params_.ap.push_back(ap);
    p_params_.app.push_back(app);
    p_params_.b.push_back(b);
    p_params_.bp.push_back(bp);
    p_params_.bpp.push_back(bpp);
    p_params_.c.push_back(c);
    p_params_.s.push_back(s);
  }
}

void Line::CalcPointScalarCovariances() {
  point_scalar_vars_.clear();
  double P;
  double inverse_P_sum = 0;
  for (std::vector<boost::array<double, 4> >::const_iterator cit =
           point_covs_.begin();
       cit != point_covs_.end(); ++cit) {
    P = (*cit)[0] * std::pow(std::cos(angle_), 2) +
        2 * (*cit)[1] * std::sin(angle_) * std::cos(angle_) +
        (*cit)[3] * std::pow(std::sin(angle_), 2);
    inverse_P_sum += 1.0 / P;
    point_scalar_vars_.push_back(P);
  }
  p_rr_ = 1.0 / inverse_P_sum;
}

void Line::RadiusFromLeastSq() {
  radius_ = 0;
  double r, phi;
  CHECK_EQ(laser_data_.size(), points_info_.size());
  for (int i = 0; i < laser_data_.size(); i++) {
    r = points_info_[i].range;
    phi = points_info_[i].angle;
    radius_ += r * std::cos(angle_ - phi) / point_scalar_vars_[i];
  }

  radius_ *= p_rr_;
}

LineExtraction::LineExtraction() {}

LineExtraction::LineExtraction(const LineExtractionParams &params)
    : params_(params) {}

LineExtraction::~LineExtraction() {}

void LineExtraction::ExtractLines(const sensor::PointCloud &laser_data,
                                  std::vector<LineInfo> *lines_info) {
  lines_info->clear();
  lines_.clear();

  FilterLaserPoints(laser_data);

  if (filtered_laser_data_.size() <= std::max(params_.min_line_points, 3)) {
    return;
  }

  // Split indices into lines and filter out short and sparse lines
  Split(filtered_laser_data_);
  FilterLines();

  // Fit each line using least squares and merge colinear lines
  for (std::vector<Line>::iterator it = lines_.begin(); it != lines_.end();
       ++it) {
    it->LeastSqFit();
  }

#if 0
  // If there is more than one line, check if lines should be merged based on
  // the merging criteria
  if (lines_.size() > 1) {
    MergeLines();
  }
#endif

  for (const auto &item : lines_) {
    LineInfo line_info;
    line_info.start = item.GetStart();
    line_info.end = item.GetEnd();
    line_info.center[0] = 0.5 * (line_info.start[0] + line_info.end[0]);
    line_info.center[1] = 0.5 * (line_info.start[1] + line_info.end[1]);
    line_info.num_points = item.NumPoints();
    line_info.length = item.Length();
    line_info.angle = item.GetAngle();
    line_info.radius = item.GetRadius();
    lines_info->emplace_back(line_info);
  }
}

double LineExtraction::DistBetweenPoints(
    const sensor::RangefinderPoint &point_1,
    const sensor::RangefinderPoint &point_2) {
  return std::sqrt(std::pow(point_1.position[0] - point_2.position[0], 2) +
                   std::pow(point_1.position[1] - point_2.position[1], 2));
}

double LineExtraction::ChiSquared(const Eigen::Vector2d &dL,
                                  const Eigen::Matrix2d &P_1,
                                  const Eigen::Matrix2d &P_2) {
  return dL.transpose() * (P_1 + P_2).inverse() * dL;
}

void LineExtraction::FilterLaserPoints(const sensor::PointCloud &laser_data) {
  filtered_laser_data_.clear();
  for (int i = 0; i < laser_data.size(); i++) {
    const sensor::RangefinderPoint &point = laser_data[i];
    const double range = std::sqrt(point.position[0] * point.position[0] +
                                   point.position[1] * point.position[1]);

    sensor::RangefinderPoint p_j, p_k;
    if (i == 0) {
      p_j = laser_data[i + 1];
      p_k = laser_data[i + 2];
    } else if (i == laser_data.size() - 1) {
      p_j = laser_data[i - 1];
      p_k = laser_data[i - 2];
    } else {
      p_j = laser_data[i - 1];
      p_k = laser_data[i + 1];
    }
    const double range_j = std::sqrt(p_j.position[0] * p_j.position[0] +
                                     p_j.position[1] * p_j.position[1]);
    const double range_k = std::sqrt(p_k.position[0] * p_k.position[0] +
                                     p_k.position[1] * p_k.position[1]);
    if (std::fabs(range - range_j) > params_.outlier_dist &&
        std::fabs(range - range_k) > params_.outlier_dist) {
      continue;
    }

    filtered_laser_data_.emplace_back(point);
  }
}

// Filtering and merging lines
void LineExtraction::FilterLines() {
  for (auto iter = lines_.begin(); iter != lines_.end();) {
    const double length = iter->Length();
    if (length < params_.min_line_length || length > params_.max_line_length ||
        iter->NumPoints() < params_.min_line_points) {
      lines_.erase(iter);
    } else {
      ++iter;
    }
  }
}

void LineExtraction::MergeLines() {
  std::vector<Line> merged_lines;

  for (std::size_t i = 1; i < lines_.size(); ++i) {
    // Get L, P_1, P_2 of consecutive lines
    Eigen::Vector2d L_1(lines_[i - 1].GetRadius(), lines_[i - 1].GetAngle());
    Eigen::Vector2d L_2(lines_[i].GetRadius(), lines_[i].GetAngle());
    Eigen::Matrix2d P_1;
    P_1 << lines_[i - 1].GetCovariance()[0], lines_[i - 1].GetCovariance()[1],
        lines_[i - 1].GetCovariance()[2], lines_[i - 1].GetCovariance()[3];
    Eigen::Matrix2d P_2;
    P_2 << lines_[i].GetCovariance()[0], lines_[i].GetCovariance()[1],
        lines_[i].GetCovariance()[2], lines_[i].GetCovariance()[3];

    // Merge lines if chi-squared distance is less than 3
    if (ChiSquared(L_1 - L_2, P_1, P_2) < 3) {
      // Get merged angle, radius, and covariance
      Eigen::Matrix2d P_m = (P_1.inverse() + P_2.inverse()).inverse();
      Eigen::Vector2d L_m = P_m * (P_1.inverse() * L_1 + P_2.inverse() * L_2);
      // Populate new line with these merged parameters
      boost::array<double, 4> cov;
      cov[0] = P_m(0, 0);
      cov[1] = P_m(0, 1);
      cov[2] = P_m(1, 0);
      cov[3] = P_m(1, 1);
      sensor::PointCloud merged_laser_data;

      const sensor::PointCloud &data_1 = lines_[i - 1].GetLaserData();
      const sensor::PointCloud &data_2 = lines_[i].GetLaserData();
      merged_laser_data.resize(data_1.size() + data_1.size());
      merged_laser_data.insert(merged_laser_data.end(), data_1.begin(),
                               data_1.end());
      merged_laser_data.insert(merged_laser_data.end(), data_2.begin(),
                               data_2.end());
      Line merged_line(merged_laser_data, params_);
      merged_line.SetLineParas(L_m[1], L_m[0], cov, lines_[i - 1].GetStart(),
                               lines_[i].GetEnd());
      // Project the new endpoints
      merged_line.ProjectEndpoints();
      lines_[i] = merged_line;
    } else {
      merged_lines.push_back(lines_[i - 1]);
    }

    if (i == lines_.size() - 1) {
      merged_lines.push_back(lines_[i]);
    }
  }
  lines_ = merged_lines;
}

// Splitting points into lines
void LineExtraction::Split(const sensor::PointCloud &laser_data) {
  if (laser_data.size() <= 3) {
    return;
  }

  Line line(laser_data, params_);
  line.EndpointFit();
  double dist_max = 0;
  double gap_max = 0;
  double dist, gap;
  int i_max, i_gap;

  for (int i = 0; i < laser_data.size() - 1; i++) {
    const sensor::RangefinderPoint &point = laser_data[i];
    dist = line.DistToPoint(point);
    if (dist > dist_max) {
      dist_max = dist;
      i_max = i;
    }
    gap = DistBetweenPoints(point, laser_data[i + 1]);
    if (gap > gap_max) {
      gap_max = gap;
      i_gap = i + 1;
    }
  }

  // Check if line meets requirements or should be split
  if (dist_max < params_.min_split_dist && gap_max < params_.max_line_gap) {
    lines_.push_back(line);
  } else {
    int i_split = dist_max >= params_.min_split_dist ? i_max : i_gap;
    sensor::PointCloud first_split(laser_data.begin(),
                                   laser_data.begin() + i_split);
    sensor::PointCloud second_split(laser_data.begin() + i_split,
                                    laser_data.end());
    Split(first_split);
    Split(second_split);
  }
}

}  // namespace mapping
}  // namespace slam
